# coding=utf-8
# Adapted from https://github.com/wzzheng/StreamVGGT/blob/main/src/eval/mv_recon/launch.py
# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
# Copyright [2025–present] StreamVGGT. All rights reserved.
# Licensed under the Creative Commons Attribution-NonCommercial-ShareAlike 4.0 License.  
#
# --------------------------------------------------------
import os
import sys
import os.path as osp
import logging

from tqdm import tqdm

import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu
from torch.utils.data._utils.collate import default_collate


sys.path.append(os.path.dirname(os.path.dirname(__file__)))
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
from dataset_utils.criterion import Regr3D_t_ScaleShiftInv, L21
from dataset_utils.eth3d import ETH3D
from vggt.models.vggt import VGGT
from vggt.utils.pose_enc import pose_encoding_to_extri_intri
from vggt.utils.geometry import unproject_depth_map_to_point_map
from vggt.utils.cast_weight import cast_model_weight
from general_utils import fix_random_seed, get_point_map_estimation_opts
from perf_metric import calc_performance
from utils import transfer_data_between_devices, \
    denormalize_image, extract_pts3d, projection_alignment, get_pcd, write_pcd
    
logging.basicConfig(level=logging.INFO)


def model_inference(model, data, dtype, use_proj):
    with torch.cuda.amp.autocast(dtype=dtype):
        with torch.no_grad():
            images = torch.stack([view['img'] for view in data], dim=0).permute(1, 0, 2, 3, 4)
            predictions = model(images)

            batch_size, seq_len = images.shape[:2]
            ress = []

            for s in range(seq_len):
                res = {
                    'pts3d_in_other_view': predictions['world_points'][:, s],  # [B, H, W, 3]
                    'conf': predictions['world_points_conf'][:, s],  # [B, H, W]

                    'depth': predictions['depth'][:, s],  # [B, H, W, 1]
                    'depth_conf': predictions['depth_conf'][:, s],  # [B, H, W]
                    'camera_pose': predictions['pose_enc'][:, s, :],  # [B, 9]
                }
                ress.append(res)
            
            preds, data = ress, data

            transfer_data_between_devices(preds, "preds", "cpu")

            if use_proj:
                pose_enc = torch.stack([preds[s]["camera_pose"] for s in range(len(preds))], dim=1)
                depth_map = torch.stack([preds[s]["depth"] for s in range(len(preds))], dim=1)
                depth_conf = torch.stack([preds[s]["depth_conf"] for s in range(len(preds))], dim=1)
                extrinsic, intrinsic = pose_encoding_to_extri_intri(pose_enc,
                                                                    data[0]["img"].shape[-2:])
                point_map_by_unprojection = unproject_depth_map_to_point_map(depth_map.squeeze(0),
                                                                                extrinsic.squeeze(0),
                                                                                intrinsic.squeeze(0))
                return preds, data, depth_conf, point_map_by_unprojection
            return preds, data, None, None


def main(model, dataset, device, args, dtype):
    name_data = "ETH3D"

    criterion = Regr3D_t_ScaleShiftInv(L21, norm_mode=False, gt_scale=True)

    with torch.no_grad():
        save_path = osp.join(args.output_dir, name_data)
        acc_all = 0
        comp_all = 0
        nc1_all = 0
        nc2_all = 0

        for data_idx in tqdm(range(len(dataset))):
            batch = default_collate([dataset[data_idx]])
            scene_id = batch[0]['label'][0].rsplit("/", 1)[0]
            transfer_data_between_devices(batch, "input", device)
            denormalize_image(batch, dtype)
            preds, batch, depth_conf, point_map_by_unprojection = \
                model_inference(model, batch, dtype, args.use_proj)
            logging.info(f"Evaluation for {name_data} {data_idx+1}/{len(dataset)}")

            transfer_data_between_devices(batch, "input", "cpu")

            pred_result = [preds, point_map_by_unprojection, depth_conf]

            pts_all_masked, pts_gt_all_masked, images_all_masked = \
               extract_pts3d(criterion, batch, pred_result, args.use_proj)
            
            if args.use_proj:
                pts_all_masked, pts_gt_all_masked, pts_all_aligned = \
                    projection_alignment(pts_all_masked, pts_gt_all_masked)

            pcd, pcd_gt = get_pcd(pts_all_masked, pts_gt_all_masked, images_all_masked)

            write_pcd(pcd, pcd_gt, save_path, scene_id)

            acc, nc1, comp, nc2 = calc_performance(pcd, pcd_gt)
            logging.info(f"Accuracy: {acc}, NC1: {nc1},  \
                    Completeness: {comp}, NC2: {nc2}")
            
            acc_all += acc
            comp_all += comp
            nc1_all += nc1
            nc2_all += nc2


            torch.cuda.empty_cache()
        acc_mean = acc_all / len(dataset)
        comp_mean = comp_all / len(dataset)
        overall_mean = (acc_mean + comp_mean) / 2
        logging.info(f"Final Results--Accuracy:{acc_mean} | Completeness: {comp_mean} | Overall: {overall_mean} ")
        

if __name__ == "__main__":
    # Parse command-line arguments
    args = get_point_map_estimation_opts()
    # Set random seeds
    fix_random_seed(42)
    # Setup device and data type
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.bfloat16 
    # Load model
    model = VGGT()
    checkpoint_path = args.ckpt  # Path to the model checkpoint
    checkpoint = torch.load(checkpoint_path)
    model.load_state_dict(checkpoint)
    model.to(device).eval()
    model = model.to(dtype)
    model = cast_model_weight(model)
    resolution = (518, 392)
    # Load dataset
    test_dataset = ETH3D(
            root_dir=args.dataset_dir,
            resolution=resolution,
            num_seq=1,
        )
    torch.npu.set_compile_mode(jit_compile=False)
    main(model, test_dataset, device, args, dtype)
    torch.cuda.empty_cache()